#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PCAP File Parser Script
PCAP文件解析脚本

功能：
1. 解析pcap文件，提取网络包信息
2. 识别TCP、UDP、ICMP、DNS等协议
3. 生成结构化的JSON数据
4. 支持批量处理

作者：Falcomm自研团队
版本：1.0
日期：2025-08-15
"""

import sys
import json
import argparse
from pathlib import Path
from datetime import datetime
from collections import defaultdict

try:
    from scapy.all import rdpcap, IP, TCP, UDP, ICMP, DNS
    SCAPY_AVAILABLE = True
except ImportError:
    SCAPY_AVAILABLE = False
    print("Warning: scapy library not installed, will use basic parsing mode")

class PCAPParser:
    def __init__(self):
        self.protocols = {
            1: "ICMP",
            6: "TCP", 
            17: "UDP",
            53: "DNS",
            132: "SCTP"  # SCTP协议，用于S1-MME接口
        }
        
        # 4G/5G核心网接口端口映射
        self.interface_ports = {
            # S1接口
            36412: "S1-MME",      # S1-MME控制面接口
            2152: "S1-U",         # S1-U用户面接口
            # S11接口
            2123: "S11",          # S11控制面接口
            # S5/S8接口
            2123: "S5/S8",        # S5/S8控制面接口
            2152: "S5/S8-U",      # S5/S8用户面接口
            # S6a接口
            3868: "S6a",          # S6a接口(Diameter)
            # N3接口
            2152: "N3",           # N3用户面接口
            # N4接口
            8805: "N4",           # N4控制面接口
            # N6接口
            2152: "N6",           # N6用户面接口
            # GTP端口
            2123: "GTP-C",        # GTP控制面
            2152: "GTP-U",        # GTP用户面
            # 其他关键端口
            5060: "SIP",          # SIP信令
            5061: "SIPS",         # 安全SIP
            16384: "RTP",         # RTP媒体流
            16385: "RTCP",        # RTCP控制
        }
        
    def parse_pcap_file(self, pcap_file_path, output_dir):
        """解析pcap文件并生成JSON"""
        pcap_file = Path(pcap_file_path)
        output_dir = Path(output_dir)
        
        if not pcap_file.exists():
            print(f"Pcap file does not exist: {pcap_file}")
            return False
            
        if not output_dir.exists():
            output_dir.mkdir(parents=True, exist_ok=True)
        
        print(f"Starting pcap file parsing: {pcap_file.name}")
        
        try:
            if SCAPY_AVAILABLE:
                packets_data = self._parse_with_scapy(pcap_file)
            else:
                packets_data = self._parse_basic(pcap_file)
            
            # 生成输出文件名
            output_file = output_dir / f"{pcap_file.stem}.json"
            
            # 保存JSON文件
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(packets_data, f, ensure_ascii=False, indent=2)
            
            print(f"JSON file generated successfully: {output_file}")
            print(f"Parsed packets: {len(packets_data['packets'])}")
            
            return True
            
        except Exception as e:
            print(f"Error parsing pcap file: {e}")
            import traceback
            traceback.print_exc()
            return False
    
    def _parse_with_scapy(self, pcap_file):
        """使用scapy解析pcap文件"""
        print("Using scapy for advanced parsing...")
        
        # 读取pcap文件
        packets = rdpcap(str(pcap_file))
        
        # 初始化数据结构
        packets_data = {
            "metadata": {
                "filename": pcap_file.name,
                "parser_version": "1.0",
                "parse_timestamp": datetime.now().isoformat(),
                "total_packets": len(packets),
                "statistics": {
                    "tcp_packets": 0,
                    "udp_packets": 0,
                    "icmp_packets": 0,
                    "dns_packets": 0,
                    "other_packets": 0
                }
            },
            "packets": []
        }
        
        # 解析每个包
        for i, packet in enumerate(packets):
            packet_info = self._extract_packet_info(packet, i + 1)
            packets_data["packets"].append(packet_info)
            
            # 更新统计信息
            self._update_statistics(packets_data["metadata"]["statistics"], packet_info)
        
        return packets_data
    
    def _parse_basic(self, pcap_file):
        """基础解析模式（当scapy不可用时）"""
        print("Using basic parsing mode...")
        
        packets_data = {
            "metadata": {
                "filename": pcap_file.name,
                "parser_version": "1.0",
                "parse_timestamp": datetime.now().isoformat(),
                "total_packets": 0,
                "note": "Basic parsing mode, limited information",
                "statistics": {
                    "tcp_packets": 0,
                    "udp_packets": 0,
                    "icmp_packets": 0,
                    "dns_packets": 0,
                    "other_packets": 0
                }
            },
            "packets": []
        }
        
        print("Basic parsing mode has limited functionality, recommend installing scapy library")
        return packets_data
    
    def _extract_packet_info(self, packet, packet_number):
        """提取单个包的信息"""
        packet_info = {
            "packet_number": packet_number,
            "timestamp": float(packet.time) if hasattr(packet, 'time') else 0,
            "datetime": datetime.fromtimestamp(float(packet.time)).isoformat() if hasattr(packet, 'time') else "",
            "length": len(packet),
            "protocol": "UNKNOWN"
        }
        
        # 提取IP层信息
        if IP in packet:
            ip_layer = packet[IP]
            packet_info.update({
                "src_ip": ip_layer.src,
                "dst_ip": ip_layer.dst,
                "protocol_num": ip_layer.proto,
                "ttl": ip_layer.ttl
            })
            
            # 识别传输层协议
            if TCP in packet:
                tcp_layer = packet[TCP]
                packet_info.update({
                    "src_port": tcp_layer.sport,
                    "dst_port": tcp_layer.dport,
                    "transport": "TCP",
                    "protocol": "TCP",
                    "tcp_flags": self._extract_tcp_flags(tcp_layer),
                    "sequence": tcp_layer.seq,
                    "acknowledgment": tcp_layer.ack,
                    "window_size": tcp_layer.window
                })
                
            elif UDP in packet:
                udp_layer = packet[UDP]
                packet_info.update({
                    "src_port": udp_layer.sport,
                    "dst_port": udp_layer.dport,
                    "transport": "UDP",
                    "protocol": "UDP"
                })
                
                # 识别4G/5G核心网接口
                interface_info = self._identify_interface(udp_layer.sport, udp_layer.dport)
                if interface_info:
                    packet_info.update(interface_info)
                
                # 检查是否为GTP协议
                if packet_info.get("protocol") in ["GTP-U", "S1AP/GTP-C"]:
                    # 提取UDP载荷进行GTP解析
                    if hasattr(udp_layer, 'payload') and udp_layer.payload:
                        gtp_info = self._extract_gtp_info(bytes(udp_layer.payload))
                        if gtp_info:
                            packet_info["gtp_info"] = gtp_info
                
                # 检查是否为DNS
                if DNS in packet:
                    packet_info["protocol"] = "DNS"
                    packet_info["dns_info"] = self._extract_dns_info(packet[DNS])
                    
            elif ICMP in packet:
                icmp_layer = packet[ICMP]
                packet_info.update({
                    "transport": "ICMP",
                    "protocol": "ICMP",
                    "icmp_type": icmp_layer.type,
                    "icmp_code": icmp_layer.code
                })
                
            else:
                # 其他IP协议
                protocol_name = self.protocols.get(ip_layer.proto, f"PROTO_{ip_layer.proto}")
                packet_info["protocol"] = protocol_name
                
                # 检查SCTP协议（S1-MME接口）
                if ip_layer.proto == 132:  # SCTP
                    packet_info["protocol"] = "SCTP"
                    packet_info["interface"] = "S1-MME"
                    packet_info["interface_port"] = 36412
                    packet_info["plane_type"] = "control"
                
        else:
            # 非IP包
            packet_info.update({
                "src_ip": "N/A",
                "dst_ip": "N/A",
                "protocol_num": 0,
                "src_port": 0,
                "dst_port": 0,
                "transport": "UNKNOWN"
            })
        
        return packet_info
    
    def _extract_tcp_flags(self, tcp_layer):
        """提取TCP标志位"""
        flags = []
        if tcp_layer.flags & 0x01:  # FIN
            flags.append("FIN")
        if tcp_layer.flags & 0x02:  # SYN
            flags.append("SYN")
        if tcp_layer.flags & 0x04:  # RST
            flags.append("RST")
        if tcp_layer.flags & 0x08:  # PSH
            flags.append("PSH")
        if tcp_layer.flags & 0x10:  # ACK
            flags.append("ACK")
        if tcp_layer.flags & 0x20:  # URG
            flags.append("URG")
        return flags
    
    def _identify_interface(self, src_port, dst_port):
        """识别4G/5G核心网接口"""
        interface_info = {}
        
        # 检查源端口和目标端口
        for port in [src_port, dst_port]:
            if port in self.interface_ports:
                interface_name = self.interface_ports[port]
                interface_info["interface"] = interface_name
                interface_info["interface_port"] = port
                
                # 根据接口类型设置协议
                if "S1-MME" in interface_name or "S11" in interface_name or "S5/S8" in interface_name:
                    interface_info["protocol"] = "S1AP/GTP-C"
                    interface_info["plane_type"] = "control"
                elif "S1-U" in interface_name or "S5/S8-U" in interface_name or "N3" in interface_name or "N6" in interface_name:
                    interface_info["protocol"] = "GTP-U"
                    interface_info["plane_type"] = "user"
                elif "S6a" in interface_name:
                    interface_info["protocol"] = "Diameter"
                    interface_info["plane_type"] = "control"
                elif "N4" in interface_name:
                    interface_info["protocol"] = "PFCP"
                    interface_info["plane_type"] = "control"
                elif "SIP" in interface_name:
                    interface_info["protocol"] = "SIP"
                    interface_info["plane_type"] = "user"
                elif "RTP" in interface_name:
                    interface_info["protocol"] = "RTP"
                    interface_info["plane_type"] = "user"
                
                break
                
        return interface_info
    
    def _extract_gtp_info(self, packet_data):
        """提取GTP协议信息"""
        gtp_info = {}
        
        if len(packet_data) >= 8:
            # GTP头部基本信息
            gtp_info["version"] = (packet_data[0] >> 5) & 0x07
            gtp_info["protocol_type"] = (packet_data[0] >> 4) & 0x01
            gtp_info["extension_header"] = (packet_data[0] >> 2) & 0x01
            gtp_info["sequence_number"] = (packet_data[0] >> 1) & 0x01
            gtp_info["n_pdu_number"] = packet_data[0] & 0x01
            gtp_info["message_type"] = packet_data[1]
            gtp_info["length"] = int.from_bytes(packet_data[2:4], 'big')
            gtp_info["teid"] = int.from_bytes(packet_data[4:8], 'big')
            
            # 根据消息类型进一步解析
            if gtp_info["message_type"] == 255:  # GTP-U
                gtp_info["gtp_type"] = "GTP-U"
                gtp_info["plane_type"] = "user"
            else:  # GTP-C
                gtp_info["gtp_type"] = "GTP-C"
                gtp_info["plane_type"] = "control"
                gtp_info["message_name"] = self._get_gtp_c_message_name(gtp_info["message_type"])
        
        return gtp_info
    
    def _get_gtp_c_message_name(self, message_type):
        """获取GTP-C消息类型名称"""
        gtp_c_messages = {
            1: "Echo Request",
            2: "Echo Response", 
            3: "Version Not Supported",
            16: "Create Session Request",
            17: "Create Session Response",
            18: "Modify Bearer Request",
            19: "Modify Bearer Response",
            20: "Delete Session Request",
            21: "Delete Session Response",
            32: "Create Bearer Request",
            33: "Create Bearer Response",
            34: "Update Bearer Request",
            35: "Update Bearer Response",
            36: "Delete Bearer Request",
            37: "Delete Bearer Response"
        }
        return gtp_c_messages.get(message_type, f"Unknown({message_type})")
    
    def _extract_dns_info(self, dns_layer):
        """提取DNS信息"""
        dns_info = {
            "qr": "Query" if dns_layer.qr == 0 else "Response",
            "opcode": dns_layer.opcode,
            "questions": dns_layer.qdcount,
            "answers": dns_layer.ancount
        }
        
        # 提取查询信息
        if hasattr(dns_layer, 'qd') and dns_layer.qd:
            dns_info["query"] = str(dns_layer.qd.qname)
            dns_info["qtype"] = dns_layer.qd.qtype
            
        return dns_info
    
    def _update_statistics(self, stats, packet_info):
        """更新统计信息"""
        protocol = packet_info.get("protocol", "UNKNOWN")
        interface = packet_info.get("interface", "")
        
        if protocol == "TCP":
            stats["tcp_packets"] += 1
        elif protocol == "UDP":
            stats["udp_packets"] += 1
        elif protocol == "ICMP":
            stats["icmp_packets"] += 1
        elif protocol == "DNS":
            stats["dns_packets"] += 1
        elif protocol == "GTP-U":
            stats["gtp_u_packets"] = stats.get("gtp_u_packets", 0) + 1
        elif protocol == "S1AP/GTP-C":
            stats["gtp_c_packets"] = stats.get("gtp_c_packets", 0) + 1
        elif protocol == "SIP":
            stats["sip_packets"] = stats.get("sip_packets", 0) + 1
        elif protocol == "RTP":
            stats["rtp_packets"] = stats.get("rtp_packets", 0) + 1
        elif protocol == "SCTP":
            stats["sctp_packets"] = stats.get("sctp_packets", 0) + 1
        else:
            stats["other_packets"] += 1
        
        # 统计接口类型
        if interface:
            if "S1-MME" in interface:
                stats["s1_mme_packets"] = stats.get("s1_mme_packets", 0) + 1
            elif "S1-U" in interface:
                stats["s1_u_packets"] = stats.get("s1_u_packets", 0) + 1
            elif "S11" in interface:
                stats["s11_packets"] = stats.get("s11_packets", 0) + 1
            elif "S5/S8" in interface:
                stats["s5_s8_packets"] = stats.get("s5_s8_packets", 0) + 1
            elif "N3" in interface:
                stats["n3_packets"] = stats.get("n3_packets", 0) + 1
            elif "N4" in interface:
                stats["n4_packets"] = stats.get("n4_packets", 0) + 1

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="PCAP File Parser Tool")
    parser.add_argument("pcap_file", help="Input pcap file path")
    parser.add_argument("output_dir", help="Output JSON file directory")
    
    args = parser.parse_args()
    
    # 创建解析器实例
    pcap_parser = PCAPParser()
    
    # 解析pcap文件
    success = pcap_parser.parse_pcap_file(args.pcap_file, args.output_dir)
    
    if success:
        print("PCAP parsing completed!")
        sys.exit(0)
    else:
        print("PCAP parsing failed!")
        sys.exit(1)

if __name__ == "__main__":
    main() 